# main.py
"""
主程序：采样、判定和数据下发脚本。

运行本脚本需要提供 YAML 配置文件。脚本将按照配置:

1. 初始化检测器，根据模型权重文件进行推理（可选）；
2. 从摄像头或视频文件读取图像，执行 YOLOv8 检测；
3. 根据中心带判定算法和多帧投票＋卡尔曼滤波得到稳定 flag；
4. 按周期采样编码器 Z 值并记录 [flag, Z] 序列；
5. 到达顶端或行程上限后，对序列进行后处理生成可执行表；
6. 打包为 Modbus TCP 格式，通过网络下发给 PLC；
7. 根据配置可选地保存采样日志或调试视频。

使用示例：

```
python -m insulator_bot.main --config config.yaml --video input.mp4
```
"""

from __future__ import annotations

import argparse
import csv
import logging
import os
import sys
import threading
import time
from typing import List

import cv2
import numpy as np
import yaml

from detection import Detector
from modbus_utils import build_payload, send_modbus_tcp
from postprocess import postprocess_sequences
from utils import KalmanFilter1D, VotingBuffer


# def judge_center_band(detections: List[List[float]], conf_thr: dict[str, float],
#                       img_height: int, band_width: int) -> int:
#     """
#     中心带可清判定函数。
#
#     若中心带内命中 class=1(片体)，则返回 1；若命中禁清类（0顶端、2法兰、3底座），优先返回 0；
#     若均未命中，则返回 0。
#
#     :param detections: 检测结果列表 [x1,y1,x2,y2,class_id,conf]
#     :param conf_thr: 各类置信度阈值字典
#     :param img_height: 图像高度
#     :param band_width: 中心带宽度（像素）
#     :return: frame_flag (0 或 1)
#     """
#     center_y = img_height // 2
#     band_half = band_width // 2
#     y_min = center_y - band_half
#     y_max = center_y + band_half
#     flag = 0
#     for x1, y1, x2, y2, cls_id, conf in detections:
#         cls = int(cls_id)
#         thr = conf_thr.get({0: 'top', 1: 'body', 2: 'flange', 3: 'base'}.get(cls, ''), 0.0)
#         if conf < thr:
#             continue
#         # 判断 bbox 是否覆盖中心线
#         if (y1-10) <= center_y <= (y2+10):
#             if cls in (0, 2, 3):
#                 return 0  # 禁清优先
#             elif cls == 1:
#                 flag = 1
#     return flag

def judge_center_band(detections: List[List[float]], conf_thr: dict[str, float],
                      img_height: int, band_width: int) -> int:
    """
    中心带可清判定函数。

    若中心带内命中 class=1(片体)，则返回 1；若命中禁清类（0顶端、2法兰、3底座），优先返回 0；
    若均未命中，则返回 0。

    :param detections: 检测结果列表 [x1,y1,x2,y2,class_id,conf]
    :param conf_thr: 各类置信度阈值字典
    :param img_height: 图像高度
    :param band_width: 中心带宽度（像素）
    :return: frame_flag (0 或 1)
    """
    # 建立类别ID到名称的映射
    class_mapping = {0: 'top', 1: 'body', 2: 'flange', 3: 'base'}

    # 计算中心带的上下边界
    center_y = img_height // 2
    band_half = band_width // 2
    band_y1 = center_y - band_half
    band_y2 = center_y + band_half

    # 优先检查是否有检测框完全包含中心带
    for x1, y1, x2, y2, cls_id, conf in detections:
        cls = int(cls_id)

        # 获取该类别的置信度阈值
        class_name = class_mapping.get(cls, '')
        thr = conf_thr.get(class_name, 0.0)

        # 置信度不足，跳过
        if conf < thr:
            continue

        # 检查检测框是否完全包含中心带
        if y1 <= band_y1 and band_y2 <= y2:
            # 完全包含，直接根据类别判断
            if cls in (0, 2, 3):  # 禁清类
                return 0
            elif cls == 1:  # 可清类（片体）
                return 1

    # 先检查禁清类
    for x1, y1, x2, y2, cls_id, conf in detections:
        cls = int(cls_id)

        # 获取该类别的置信度阈值
        class_name = class_mapping.get(cls, '')
        thr = conf_thr.get(class_name, 0.0)

        # 置信度不足，跳过
        if conf < thr:
            continue

        # 计算检测框与中心带的重叠区域
        overlap_y1 = max(band_y1, y1)
        overlap_y2 = min(band_y2, y2)

        # 如果有重叠区域
        if overlap_y1 < overlap_y2:
            overlap_height = overlap_y2 - overlap_y1
            box_height = y2 - y1

            # 计算重叠比例（相对于检测框高度）
            if box_height > 0:
                overlap_ratio = overlap_height / box_height
                # 如果重叠比例超过50%，且是禁清类
                if overlap_ratio > 0.5 and cls in (0, 2, 3):
                    return 0

    # 再检查可清类（片体）
    for x1, y1, x2, y2, cls_id, conf in detections:
        cls = int(cls_id)

        # 获取该类别的置信度阈值
        class_name = class_mapping.get(cls, '')
        thr = conf_thr.get(class_name, 0.0)

        # 置信度不足，跳过
        if conf < thr:
            continue

        # 计算检测框与中心带的重叠区域
        overlap_y1 = max(band_y1, y1)
        overlap_y2 = min(band_y2, y2)

        # 如果有重叠区域
        if overlap_y1 < overlap_y2:
            overlap_height = overlap_y2 - overlap_y1
            box_height = y2 - y1

            # 计算重叠比例（相对于检测框高度）
            if box_height > 0:
                overlap_ratio = overlap_height / box_height
                # 如果重叠比例超过50%，且是片体
                if overlap_ratio > 0.5 and cls == 1:
                    return 1

    return 0  # 默认返回禁清

    return 0  # 默认返回禁清


class EncoderSimulator:
    """
    简易编码器模拟器，用于在无真实硬件时生成单调递增的 Z 值。

    根据设置的最高行程 ``z_max`` 和期望持续时间 ``duration``（秒），按时间推进产生 Z。
    """

    def __init__(self, z_max: float, duration: float = 30.0) -> None:
        self.z_max = z_max
        self.duration = duration
        self.start_time = time.time()

    # TODO：change height
    def read(self) -> float:
        elapsed = time.time() - self.start_time
        # 按线性速度计算位置
        pos = (elapsed / self.duration) * self.z_max*0.1
        return min(pos, self.z_max)


def read_distance() -> float:
    """读测距仪距离。实际项目应从硬件读取，此处返回 -1 表示未知。"""
    return -1.0


def save_csv(csv_path: str, records: List[List[float]]) -> None:
    """保存采样序列到 CSV 文件。"""
    os.makedirs(os.path.dirname(csv_path), exist_ok=True)
    with open(csv_path, mode='w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['flag', 'z'])
        writer.writerows(records)


def main(config_path: str, video_path: str | None = None) -> None:
    # 读取 YAML 配置
    with open(config_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    vision_cfg = config.get('vision', {})
    sampling_cfg = config.get('sampling', {})
    post_cfg = config.get('postproc', {})
    modbus_cfg = config.get('modbus', {})
    logging_cfg = config.get('logging', {})
    # 初始化检测器
    weight_path = vision_cfg.get('weight_path')
    detector = Detector(weight_path)  # 如果未提供权重，则使用 DummyDetector
    # 中心带参数
    center_band_px = int(vision_cfg.get('center_band_px', 20))
    vote_k = int(vision_cfg.get('vote_k', 5))
    vote_t = int(vision_cfg.get('vote_t', 3))
    kalman_q = float(vision_cfg.get('kalman_q', 1e-3))
    kalman_r = float(vision_cfg.get('kalman_r', 5e-2))
    conf_thr = vision_cfg.get('conf_thr', {
        'tip': 0.28,
        'disk': 0.40,
        'flange': 0.55,
        'base': 0.30
    })
    # 采样参数
    period_s = float(sampling_cfg.get('period_s', 1.0))
    stop_on_tip = bool(sampling_cfg.get('stop_on_tip', True))
    z_max = float(sampling_cfg.get('z_max', 2500.0))
    # 后处理参数
    open_close_win = int(post_cfg.get('open_close_win', 3))
    min_segment_mm = float(post_cfg.get('min_segment_mm', 60))
    safety_delta_mm = float(post_cfg.get('safety_delta_mm', 25))
    brush_offset_mm = float(post_cfg.get('brush_offset_mm', 0))
    merge_gap_mm = float(post_cfg.get('merge_gap_mm', 30))
    output_mode = post_cfg.get('output_mode', 'segments')
    # Modbus 参数
    host = modbus_cfg.get('host', '127.0.0.1')
    port = int(modbus_cfg.get('port', 502))
    unit_id = int(modbus_cfg.get('unit_id', 1))
    reg_base = int(modbus_cfg.get('reg_base', 0))
    # 日志参数
    save_video = bool(logging_cfg.get('save_video', False))
    csv_path = logging_cfg.get('csv_path', None)
    # 初始化滤波器与投票器
    kalman = KalmanFilter1D(initial_state=0.0, initial_variance=1.0, q=kalman_q, r=kalman_r)
    voter = VotingBuffer(window_size=vote_k, vote_threshold=vote_t)
    # 初始化编码器（真实环境应替换为硬件读取）
    encoder = EncoderSimulator(z_max=z_max, duration=30.0)
    # 采样数据
    sample_flags: List[int] = []
    sample_zs: List[float] = []
    # sample_class:List[str] = ['empty','empty']
    running = True
    # 采样和检测循环
    last_sample_time = time.time()
    tip_detected = False
    # 视频捕获
    cap = None
    if video_path:
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            logging.error("无法打开视频文件 %s", video_path)
            cap = None

    logging.info("开始采样...")
    while running:
        # 读取帧
        if cap:
            ret, frame = cap.read()
            if not ret:
                logging.info("视频播放结束，停止采样。")
                running = False
                break
            # 将960*720的视频帧压缩为640*480
            frame = cv2.resize(frame, (640, 480))
        else:
            # 若没有视频，构造黑图
            frame = np.zeros((640, 480, 3), dtype=np.uint8)
        img_h, img_w = frame.shape[:2]
        detections = detector.detect(frame)
        # 是否处于中间带
        frame_flag = judge_center_band(detections, conf_thr, img_h, center_band_px)
        # 多帧投票
        stable_flag = voter.update(frame_flag)
        # 卡尔曼滤波
        prob = kalman.update(float(frame_flag))

        # TODO：第二阶段交互：1.检测到顶端。发送[flag_top]
        # # 提前检测顶端 class=0
        # if stop_on_tip:
        #     # 如果检测结果中存在满足阈值的 class 0 且稳定，则置为 True
        #     class_=[]
        #     for x1, y1, x2, y2, cls_id, conf in detections:
        #         class_.append(cls_id)
        #         if int(cls_id) == 0 and conf >= conf_thr.get('tip', 0.28):
        #
        #             tip_detected = True
        #             break
        #     logging.info("检测到的类别：%s",','.join( class_) if class_ else '无')
        # 提前检测顶端 class=0
        if stop_on_tip:
            # 如果检测结果中存在满足阈值的 class 0 且稳定，则置为 True
            class_names = []
            for x1, y1, x2, y2, cls_id, conf in detections:
                # 映射类别ID到名称
                class_name = {0: 'top', 1: 'body', 2: 'flange', 3: 'base'}.get(int(cls_id), 'unknown')
                class_names.append(f"{class_name}({int(cls_id)})")
                if int(cls_id) == 0 and conf >= conf_thr.get('tip', 0.28):
                    tip_detected = True
                    break
            # 输出检测到的类别名称和ID
            logging.info("检测到的类别：%s", ', '.join(class_names) if class_names else '无')
        # 周期采样
        now = time.time()
        if now - last_sample_time >= period_s:



            last_sample_time = now
            # 采样编码器值        TODO：第一阶段 发送时间戳/，通过通讯协议获得高度信息
            #                   TODO：发送[时间戳]  / 返回[绝对高度]
            z = encoder.read()
            sample_flags.append(int(stable_flag))
            sample_zs.append(float(z))
            # sample_class.append(cls_id)
            logging.info("采样 flag=%d, z=%.2f", stable_flag, z)
            # 判断结束条件
            # if stop_on_tip and tip_detected:
            #     logging.info("检测到顶端，提前结束采样。")
            #     running = False

            # TODO: 第二阶段2.到达行程上限 收到[flag_top]
            if z >= z_max:
                logging.info("达到行程上限 %.2f，结束采样。", z_max)
                running = False
    # 采样结束
    if cap:
        cap.release()
    # 后处理
    logging.info("后处理采样序列...")
    processed = postprocess_sequences(
        flags=sample_flags,
        zs=sample_zs,
        open_close_win=open_close_win,
        min_segment_mm=min_segment_mm,
        safety_delta_mm=safety_delta_mm,
        brush_offset_mm=brush_offset_mm,
        merge_gap_mm=merge_gap_mm,
        output_mode=output_mode,
    )
    # print("后处理数据", processed)
    # logging.info("后处理数据: ",processed)


    # TODO：第三阶段 逐段发送[Z_x,delta_s,n,dis,flag_last]
    # 打包并发送
    logging.info("打包并发送 Modbus TCP 数据...")
    distance = read_distance()
    header_regs, data_regs = build_payload(processed, distance, mode=output_mode)
    try:
        send_modbus_tcp(host, port, unit_id, reg_base, header_regs, data_regs)
        logging.info("Modbus 数据已发送。")
    except Exception as e:
        logging.error("发送 Modbus 数据失败：%s", e)
    # 保存 CSV
    if csv_path:
        logging.info("保存 CSV 日志到 %s", csv_path)
        save_csv(csv_path, [[f, z] for f, z in zip(sample_flags, sample_zs)])
    logging.info("采样处理完成。")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="绝缘子清洗机器人采样判定脚本")
    parser.add_argument('--config', required=True, help='YAML 配置文件路径')
    parser.add_argument('--video',default="video/demo2-1.mp4", help='可选，输入视频文件路径')
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s: %(message)s')
    main(args.config, args.video)